package com.niit.offline.rec

import breeze.numerics.sqrt
import org.apache.spark.SparkConf
import org.apache.spark.mllib.recommendation.{ALS, MatrixFactorizationModel, Rating}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession

object ALSTrainer {

  def main(args: Array[String]): Unit = {
    // 配置mongo连接信息
    val conf = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://fooadmin:123456@127.0.0.1:27017/bigdata",
      "mongo.db" -> "bigdata"
    )
    // 创建sparkConf对象
    val sparkConf = new SparkConf().setAppName("ALSTrainer").setMaster(conf("spark.cores"))
    // 创建SparkSession对象
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    // 创建mongoConfig样例类对象
    val mongoConfig = MongoConfig(conf("mongo.uri"),conf("mongo.db"))
    // 引入sparksql需要的隐式类
    import spark.implicits._
    // 读取mongo中的数据
    val ratingRDD = spark.read
      .option("uri", mongoConfig.uri)
      .option("collection", OfflineRecommender.MONGODB_RATING_COLLECTION)
      .format("com.mongodb.spark.sql")
      .load()
      .as[ProductRating]
      .rdd
      .map(
        rating => Rating(rating.userId,rating.productId,rating.score)
      ).cache() // 为了提高执行效率先做些缓存
    // 将ratingDF切分成训练集与测试集
    //随机切分 80%的数据作为训练集，20%作为测试集
    val splits = ratingRDD.randomSplit(Array(0.8,0.2))
    val trainRDD = splits(0) // 训练集
    val testRDD = splits(1) // 测试集

    // 核心代码：输出最优参数，即找到rmse最小值
    adjustALSParams( trainRDD, testRDD)
    // 关闭spark
    spark.stop()
  }
  // 定义函数：adjustALSParams，找出最优的rmse
  def adjustALSParams(trainRDD: RDD[Rating], testRDD: RDD[Rating]): Unit = {
    // 根据(rank,iterations,lambda)参数，求出模型中的最优的rmse
    val result = for ( rank <- Array(5,10,20,50); lambda <- Array(1,0.1,0.01))
      // 针对每一次循环生成的结果，需要被记录下来，故使用yield来生成一个新的集合并记录每次的结果
      yield {
        // 使用每次生成的rank和lambda作为模型参数进行模型训练
        val model = ALS.train(trainRDD,rank,10,lambda)
        val rmse = getRMSE( model,testRDD )  // 将训练得到的模型给定测试集进行评估
        (rank,lambda,rmse) // 将得到的rmse与之对应的rank,lambda组成三元组返回
      }
    // 打印结果,取出前几个
    println(result.sortBy(_._3).head)
  }
  // 编码实现求rmse公式：
  def getRMSE(model: MatrixFactorizationModel, testData: RDD[Rating]):Double = {
    // 1.得到model预测时所需要的userProduces测试数据：即从测试集数据中得到
    val userProducts = testData.map(item => (item.user,item.product))
    // 2.通过模型计算出预测数据集(predicted)；
    val predictedRdd = model.predict(userProducts)
    // 3.通过testData测试数据集得出用户对商品打分的真实值，数据结构：((userId,productId),rating)
    val observed = testData.map(item => ((item.user,item.product),item.rating))
    // 4.得到预测值
   val predict = predictedRdd.map(item => ((item.user, item.product), item.rating))
    // 5.组织成公式进行计算,
    sqrt( // Returns the square root of a `double` value.
      observed.join(predict) //通过join将真实数据集与预测数据进行内连接组成一个新的集合
        .map{
          case ((userId,productId),(real,predict)) =>
            val error = real - predict
            error * error
      }.mean() // mean计算RDD所有元素的平均值
    )
  }
}
